{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "V2BNNeQQ3pLa"
      },
      "outputs": [],
      "source": [
        "# Based on https://github.com/blakechi/ComVEX/blob/master/comvex/coatnet/model.py\n",
        "\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "from MBConv import MBConvForRelativeAttention\n",
        "from Transformer import TransformerWithRelativeAttention, ProjectionHead\n",
        "\n",
        "configs = {\n",
        "    'coatnet-0': {\n",
        "        'num_blocks': [2, 2, 3, 5, 2],\n",
        "        'num_channels': [64, 96, 192, 384, 768],\n",
        "        'expand_ratio': [4, 4, 4, 4, 4],\n",
        "        'n_head': 32,\n",
        "        'block_types': ['C', 'C', 'T', 'T']\n",
        "    },\n",
        "    'coatnet-1': {\n",
        "        'num_blocks': [2, 2, 6, 14, 2],\n",
        "        'num_channels': [64, 96, 192, 384, 768],\n",
        "        'expand_ratio': [4, 4, 4, 4, 4],\n",
        "        'n_head': 32,\n",
        "        'block_types': ['C', 'C', 'T', 'T']\n",
        "    },\n",
        "    'coatnet-2': {\n",
        "        'num_blocks': [2, 2, 6, 14, 2],\n",
        "        'num_channels': [128, 128, 256, 512, 1026],\n",
        "        'expand_ratio': [4, 4, 4, 4, 4],\n",
        "        'n_head': 32,\n",
        "        'block_types': ['C', 'C', 'T', 'T']\n",
        "    },\n",
        "    'coatnet-3': {\n",
        "        'num_blocks': [2, 2, 6, 14, 2],\n",
        "        'num_channels': [192, 192, 384, 768, 1536],\n",
        "        'expand_ratio': [4, 4, 4, 4, 4],\n",
        "        'n_head': 32,\n",
        "        'block_types': ['C', 'C', 'T', 'T']\n",
        "    },\n",
        "    'coatnet-4': {\n",
        "        'num_blocks': [2, 2, 12, 28, 2],\n",
        "        'num_channels': [192, 192, 384, 768, 1536],\n",
        "        'expand_ratio': [4, 4, 4, 4, 4],\n",
        "        'n_head': 32,\n",
        "        'block_types': ['C', 'C', 'T', 'T']\n",
        "    },\n",
        "    'coatnet-5': {\n",
        "        'num_blocks': [2, 2, 12, 28, 2],\n",
        "        'num_channels': [192, 256, 512, 1280, 2048],\n",
        "        'expand_ratio': [4, 4, 4, 4, 4],\n",
        "        'n_head': 64,\n",
        "        'block_types': ['C', 'C', 'T', 'T']\n",
        "    },\n",
        "    # Something's not right with this one\n",
        "    # 'coatnet-6': {\n",
        "    #     'num_blocks': [2, 2, 4, [8, 42], 2],\n",
        "    #     'num_channels': [192, 192, 384, [768, 1536], 2048],\n",
        "    #     'expand_ratio': [4, 4, 4, 4, 4],\n",
        "    #     'n_head': 128,\n",
        "    #     'block_types': ['C', 'C', 'C-T', 'T']\n",
        "    # },\n",
        "    # 'coatnet-7': {\n",
        "    #     'num_blocks': [2, 2, 4, [8, 42], 2],\n",
        "    #     'num_channels': [192, 256, 512, [1024, 2048], 3072],\n",
        "    #     'expand_ratio': [4, 4, 4, 4, 4],\n",
        "    #     'n_head': 128,\n",
        "    #     'block_types': ['C', 'C', 'C-T', 'T']\n",
        "    # }\n",
        "}\n",
        "\n",
        "blocks = {\n",
        "    'C': MBConvForRelativeAttention,\n",
        "    'T': TransformerWithRelativeAttention\n",
        "}\n",
        "\n",
        "\n",
        "class CoAtNet(nn.Module):\n",
        "    def __init__(self, inp_h, inp_w, in_channels, config='coatnet-0', num_classes=None, head_act_fn='mish',\n",
        "                 head_dropout=0.1):\n",
        "        super().__init__()\n",
        "        self.config = configs[config]\n",
        "        block_types = self.config['block_types']\n",
        "        self.s0 = self._make_stem(in_channels)\n",
        "        self.s1 = self._make_block(block_types[0], inp_h >> 2, inp_w >> 2,\n",
        "                                   self.config['num_channels'][0],\n",
        "                                   self.config['num_channels'][1],\n",
        "                                   self.config['num_blocks'][1],\n",
        "                                   self.config['expand_ratio'][0])\n",
        "        self.s2 = self._make_block(block_types[1], inp_h >> 3, inp_w >> 3,\n",
        "                                   self.config['num_channels'][1],\n",
        "                                   self.config['num_channels'][2],\n",
        "                                   self.config['num_blocks'][2],\n",
        "                                   self.config['expand_ratio'][1])\n",
        "        self.s3 = self._make_block(block_types[2], inp_h >> 4, inp_w >> 4,\n",
        "                                   self.config['num_channels'][2],\n",
        "                                   self.config['num_channels'][3],\n",
        "                                   self.config['num_blocks'][3],\n",
        "                                   self.config['expand_ratio'][2])\n",
        "        self.s4 = self._make_block(block_types[3], inp_h >> 5, inp_w >> 5,\n",
        "                                   self.config['num_channels'][3],\n",
        "                                   self.config['num_channels'][4],\n",
        "                                   self.config['num_blocks'][4],\n",
        "                                   self.config['expand_ratio'][3])\n",
        "        self.include_head = num_classes is not None\n",
        "        if self.include_head:\n",
        "            if isinstance(num_classes, int):\n",
        "                self.single_head = True\n",
        "                num_classes = [num_classes]\n",
        "            else:\n",
        "                self.single_head = False\n",
        "            self.heads = nn.ModuleList(\n",
        "                [ProjectionHead(self.config['num_channels'][-1], nc, act_fn=head_act_fn, ff_dropout=head_dropout) for nc\n",
        "                 in num_classes])\n",
        "\n",
        "    def _make_stem(self, in_channels):\n",
        "        return nn.Sequential(*[\n",
        "            nn.Conv2d(\n",
        "                in_channels if i == 0 else self.config['num_channels'][0],\n",
        "                self.config['num_channels'][0], kernel_size=3, padding=1,\n",
        "                stride=2 if i == 0 else 1\n",
        "            ) for i in range(self.config['num_blocks'][0])\n",
        "        ])\n",
        "\n",
        "    # 'coatnet-0': {\n",
        "    #     'num_blocks': [2, 2, 3, 5, 2],\n",
        "    #     'num_channels': [64, 96, 192, 384, 768],\n",
        "    #     'expand_ratio': [4, 4, 4, 4, 4],\n",
        "    #     'n_head': 32,\n",
        "    #     'block_types': ['C', 'C', 'T', 'T']\n",
        "    # },\n",
        "\n",
        "    def _make_block(self, block_type, inp_h, inp_w, in_channels, out_channels, depth, expand_ratio):\n",
        "        block_list = []\n",
        "        if not isinstance(in_channels, int):\n",
        "            in_channels = in_channels[-1]\n",
        "        if block_type in blocks:\n",
        "            block_cls = blocks[block_type]\n",
        "            block_list.extend([\n",
        "                block_cls(\n",
        "                    inp_h, inp_w, in_channels if i == 0 else out_channels,\n",
        "                    n_head=self.config['n_head'], out_channels=out_channels,\n",
        "                    expand_ratio=expand_ratio, use_downsampling=(i == 0)\n",
        "                ) for i in range(depth)\n",
        "            ])\n",
        "        else:\n",
        "            for i, _block_type in enumerate(block_type.split('-')):\n",
        "                block_cls = blocks[_block_type]\n",
        "                block_list.extend(\n",
        "                    block_cls(\n",
        "                        inp_h, inp_w,\n",
        "                        in_channels if i == 0 and j == 0 else out_channels[i - 1] if j == 0 else out_channels[i],\n",
        "                        n_head=self.config['n_head'], out_channels=out_channels[i],\n",
        "                        expand_ratio=expand_ratio, use_downsampling=j == 0\n",
        "                    ) for j in range(depth[i])\n",
        "                )\n",
        "        return nn.Sequential(*block_list)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = self.s0(x)\n",
        "        x = self.s1(x)\n",
        "        x = self.s2(x)\n",
        "        x = self.s3(x)\n",
        "        x = self.s4(x)\n",
        "        x = F.adaptive_avg_pool2d(x, 1).view(x.size(0), -1)\n",
        "        if self.include_head:\n",
        "            if self.single_head:\n",
        "                return self.heads[0](x)\n",
        "            return [head(x) for head in self.heads]\n",
        "        return x\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "S1Oatdum31Md",
        "outputId": "dd174ef7-604a-4bb6-fabc-465f73ce5275"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Using device: cpu\n"
          ]
        }
      ],
      "source": [
        "import os\n",
        "import sys\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.optim as optim\n",
        "import torchvision\n",
        "import torchvision.transforms as transforms\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "print(f\"Using device: {device}\")\n",
        "\n",
        "\n",
        "batch_size = 8\n",
        "num_epochs = 3\n",
        "learning_rate = 0.001\n",
        "num_classes = 100\n",
        "model_config = 'coatnet-0'\n",
        "\n",
        "# 数据预处理和加载\n",
        "transform = transforms.Compose([\n",
        "    transforms.Resize((224, 224)), \n",
        "    transforms.ToTensor(),\n",
        "])\n",
        "\n",
        "trainset = torchvision.datasets.CIFAR100(root='../data', train=True, download=True, transform=transform)\n",
        "testset = torchvision.datasets.CIFAR100(root='../data', train=False, download=True, transform=transform)\n",
        "\n",
        "trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, shuffle=True, num_workers=2)\n",
        "testloader = torch.utils.data.DataLoader(testset, batch_size=batch_size, shuffle=False, num_workers=2)\n",
        "\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "0RcYuHkf34K-",
        "outputId": "60719b05-5b48-4864-8774-6110b2ebd6e4"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "CoAtNet(\n",
            "  (s0): Sequential(\n",
            "    (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
            "    (1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
            "  )\n",
            "  (s1): Sequential(\n",
            "    (0): MBConvForRelativeAttention(\n",
            "      (norm): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
            "      (mbconv): MBConv(\n",
            "        (act_fn): Mish()\n",
            "        (expand_conv): Conv2dBlock(\n",
            "          (conv): Conv2dStaticSamePadding(\n",
            "            64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
            "            (static_padding): Identity()\n",
            "          )\n",
            "          (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
            "          (act_fn): Mish()\n",
            "        )\n",
            "        (depthwise_conv): DepthwiseConv2dBlock(\n",
            "          (conv): Conv2dStaticSamePadding(\n",
            "            256, 256, kernel_size=(3, 3), stride=(2, 2), groups=256, bias=False\n",
            "            (static_padding): ZeroPad2d((np.int64(0), np.int64(1), np.int64(0), np.int64(1)))\n",
            "          )\n",
            "          (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
            "          (act_fn): Mish()\n",
            "        )\n",
            "        (se): SqueezeExcitation(\n",
            "          (reduce): Conv2dStaticSamePadding(\n",
            "            256, 64, kernel_size=(1, 1), stride=(1, 1)\n",
            "            (static_padding): Identity()\n",
            "          )\n",
            "          (act_fn): Mish()\n",
            "          (expand): Conv2dStaticSamePadding(\n",
            "            64, 256, kernel_size=(1, 1), stride=(1, 1)\n",
            "            (static_padding): Identity()\n",
            "          )\n",
            "        )\n",
            "        (project_conv): Conv2dBlock(\n",
            "          (conv): Conv2dStaticSamePadding(\n",
            "            256, 96, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
            "            (static_padding): Identity()\n",
            "          )\n",
            "          (bn): BatchNorm2d(96, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
            "        )\n",
            "      )\n",
            "      (pool): MaxPool2d(kernel_size=(2, 2), stride=(2, 2), padding=0, dilation=1, ceil_mode=False)\n",
            "      (conv): Conv2d(64, 96, kernel_size=(1, 1), stride=(1, 1))\n",
            "    )\n",
            "    (1): MBConvForRelativeAttention(\n",
            "      (norm): BatchNorm2d(96, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
            "      (mbconv): MBConv(\n",
            "        (act_fn): Mish()\n",
            "        (expand_conv): Conv2dBlock(\n",
            "          (conv): Conv2dStaticSamePadding(\n",
            "            96, 384, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
            "            (static_padding): Identity()\n",
            "          )\n",
            "          (bn): BatchNorm2d(384, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
            "          (act_fn): Mish()\n",
            "        )\n",
            "        (depthwise_conv): DepthwiseConv2dBlock(\n",
            "          (conv): Conv2dStaticSamePadding(\n",
            "            384, 384, kernel_size=(3, 3), stride=(1, 1), groups=384, bias=False\n",
            "            (static_padding): ZeroPad2d((np.int64(1), np.int64(1), np.int64(1), np.int64(1)))\n",
            "          )\n",
            "          (bn): BatchNorm2d(384, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
            "          (act_fn): Mish()\n",
            "        )\n",
            "        (se): SqueezeExcitation(\n",
            "          (reduce): Conv2dStaticSamePadding(\n",
            "            384, 96, kernel_size=(1, 1), stride=(1, 1)\n",
            "            (static_padding): Identity()\n",
            "          )\n",
            "          (act_fn): Mish()\n",
            "          (expand): Conv2dStaticSamePadding(\n",
            "            96, 384, kernel_size=(1, 1), stride=(1, 1)\n",
            "            (static_padding): Identity()\n",
            "          )\n",
            "        )\n",
            "        (project_conv): Conv2dBlock(\n",
            "          (conv): Conv2dStaticSamePadding(\n",
            "            384, 96, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
            "            (static_padding): Identity()\n",
            "          )\n",
            "          (bn): BatchNorm2d(96, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
            "        )\n",
            "      )\n",
            "    )\n",
            "  )\n",
            "  (s2): Sequential(\n",
            "    (0): MBConvForRelativeAttention(\n",
            "      (norm): BatchNorm2d(96, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
            "      (mbconv): MBConv(\n",
            "        (act_fn): Mish()\n",
            "        (expand_conv): Conv2dBlock(\n",
            "          (conv): Conv2dStaticSamePadding(\n",
            "            96, 384, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
            "            (static_padding): Identity()\n",
            "          )\n",
            "          (bn): BatchNorm2d(384, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
            "          (act_fn): Mish()\n",
            "        )\n",
            "        (depthwise_conv): DepthwiseConv2dBlock(\n",
            "          (conv): Conv2dStaticSamePadding(\n",
            "            384, 384, kernel_size=(3, 3), stride=(2, 2), groups=384, bias=False\n",
            "            (static_padding): ZeroPad2d((np.int64(0), np.int64(1), np.int64(0), np.int64(1)))\n",
            "          )\n",
            "          (bn): BatchNorm2d(384, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
            "          (act_fn): Mish()\n",
            "        )\n",
            "        (se): SqueezeExcitation(\n",
            "          (reduce): Conv2dStaticSamePadding(\n",
            "            384, 96, kernel_size=(1, 1), stride=(1, 1)\n",
            "            (static_padding): Identity()\n",
            "          )\n",
            "          (act_fn): Mish()\n",
            "          (expand): Conv2dStaticSamePadding(\n",
            "            96, 384, kernel_size=(1, 1), stride=(1, 1)\n",
            "            (static_padding): Identity()\n",
            "          )\n",
            "        )\n",
            "        (project_conv): Conv2dBlock(\n",
            "          (conv): Conv2dStaticSamePadding(\n",
            "            384, 192, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
            "            (static_padding): Identity()\n",
            "          )\n",
            "          (bn): BatchNorm2d(192, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
            "        )\n",
            "      )\n",
            "      (pool): MaxPool2d(kernel_size=(2, 2), stride=(2, 2), padding=0, dilation=1, ceil_mode=False)\n",
            "      (conv): Conv2d(96, 192, kernel_size=(1, 1), stride=(1, 1))\n",
            "    )\n",
            "    (1): MBConvForRelativeAttention(\n",
            "      (norm): BatchNorm2d(192, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
            "      (mbconv): MBConv(\n",
            "        (act_fn): Mish()\n",
            "        (expand_conv): Conv2dBlock(\n",
            "          (conv): Conv2dStaticSamePadding(\n",
            "            192, 768, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
            "            (static_padding): Identity()\n",
            "          )\n",
            "          (bn): BatchNorm2d(768, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
            "          (act_fn): Mish()\n",
            "        )\n",
            "        (depthwise_conv): DepthwiseConv2dBlock(\n",
            "          (conv): Conv2dStaticSamePadding(\n",
            "            768, 768, kernel_size=(3, 3), stride=(1, 1), groups=768, bias=False\n",
            "            (static_padding): ZeroPad2d((np.int64(1), np.int64(1), np.int64(1), np.int64(1)))\n",
            "          )\n",
            "          (bn): BatchNorm2d(768, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
            "          (act_fn): Mish()\n",
            "        )\n",
            "        (se): SqueezeExcitation(\n",
            "          (reduce): Conv2dStaticSamePadding(\n",
            "            768, 192, kernel_size=(1, 1), stride=(1, 1)\n",
            "            (static_padding): Identity()\n",
            "          )\n",
            "          (act_fn): Mish()\n",
            "          (expand): Conv2dStaticSamePadding(\n",
            "            192, 768, kernel_size=(1, 1), stride=(1, 1)\n",
            "            (static_padding): Identity()\n",
            "          )\n",
            "        )\n",
            "        (project_conv): Conv2dBlock(\n",
            "          (conv): Conv2dStaticSamePadding(\n",
            "            768, 192, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
            "            (static_padding): Identity()\n",
            "          )\n",
            "          (bn): BatchNorm2d(192, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
            "        )\n",
            "      )\n",
            "    )\n",
            "    (2): MBConvForRelativeAttention(\n",
            "      (norm): BatchNorm2d(192, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
            "      (mbconv): MBConv(\n",
            "        (act_fn): Mish()\n",
            "        (expand_conv): Conv2dBlock(\n",
            "          (conv): Conv2dStaticSamePadding(\n",
            "            192, 768, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
            "            (static_padding): Identity()\n",
            "          )\n",
            "          (bn): BatchNorm2d(768, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
            "          (act_fn): Mish()\n",
            "        )\n",
            "        (depthwise_conv): DepthwiseConv2dBlock(\n",
            "          (conv): Conv2dStaticSamePadding(\n",
            "            768, 768, kernel_size=(3, 3), stride=(1, 1), groups=768, bias=False\n",
            "            (static_padding): ZeroPad2d((np.int64(1), np.int64(1), np.int64(1), np.int64(1)))\n",
            "          )\n",
            "          (bn): BatchNorm2d(768, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
            "          (act_fn): Mish()\n",
            "        )\n",
            "        (se): SqueezeExcitation(\n",
            "          (reduce): Conv2dStaticSamePadding(\n",
            "            768, 192, kernel_size=(1, 1), stride=(1, 1)\n",
            "            (static_padding): Identity()\n",
            "          )\n",
            "          (act_fn): Mish()\n",
            "          (expand): Conv2dStaticSamePadding(\n",
            "            192, 768, kernel_size=(1, 1), stride=(1, 1)\n",
            "            (static_padding): Identity()\n",
            "          )\n",
            "        )\n",
            "        (project_conv): Conv2dBlock(\n",
            "          (conv): Conv2dStaticSamePadding(\n",
            "            768, 192, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
            "            (static_padding): Identity()\n",
            "          )\n",
            "          (bn): BatchNorm2d(192, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
            "        )\n",
            "      )\n",
            "    )\n",
            "  )\n",
            "  (s3): Sequential(\n",
            "    (0): TransformerWithRelativeAttention(\n",
            "      (pool): MaxPool2d(kernel_size=(2, 2), stride=(2, 2), padding=0, dilation=1, ceil_mode=False)\n",
            "      (conv): Conv2d(192, 384, kernel_size=(1, 1), stride=(1, 1))\n",
            "      (norm): LayerNorm((192,), eps=1e-05, elementwise_affine=True)\n",
            "      (attention): RelativeAttention(\n",
            "        (Q): Linear(in_features=192, out_features=384, bias=False)\n",
            "        (K): Linear(in_features=192, out_features=384, bias=False)\n",
            "        (V): Linear(in_features=192, out_features=384, bias=False)\n",
            "        (ff): Linear(in_features=384, out_features=384, bias=True)\n",
            "        (attn_dropout): Dropout2d(p=0.1, inplace=False)\n",
            "        (ff_dropout): Dropout(p=0.1, inplace=False)\n",
            "      )\n",
            "      (ff): FeedForwardRelativeAttention(\n",
            "        (fc1): Conv2d(384, 1536, kernel_size=(1, 1), stride=(1, 1))\n",
            "        (fc2): Conv2d(1536, 384, kernel_size=(1, 1), stride=(1, 1))\n",
            "      )\n",
            "    )\n",
            "    (1): TransformerWithRelativeAttention(\n",
            "      (norm): LayerNorm((384,), eps=1e-05, elementwise_affine=True)\n",
            "      (attention): RelativeAttention(\n",
            "        (Q): Linear(in_features=384, out_features=384, bias=False)\n",
            "        (K): Linear(in_features=384, out_features=384, bias=False)\n",
            "        (V): Linear(in_features=384, out_features=384, bias=False)\n",
            "        (ff): Linear(in_features=384, out_features=384, bias=True)\n",
            "        (attn_dropout): Dropout2d(p=0.1, inplace=False)\n",
            "        (ff_dropout): Dropout(p=0.1, inplace=False)\n",
            "      )\n",
            "      (ff): FeedForwardRelativeAttention(\n",
            "        (fc1): Conv2d(384, 1536, kernel_size=(1, 1), stride=(1, 1))\n",
            "        (fc2): Conv2d(1536, 384, kernel_size=(1, 1), stride=(1, 1))\n",
            "      )\n",
            "    )\n",
            "    (2): TransformerWithRelativeAttention(\n",
            "      (norm): LayerNorm((384,), eps=1e-05, elementwise_affine=True)\n",
            "      (attention): RelativeAttention(\n",
            "        (Q): Linear(in_features=384, out_features=384, bias=False)\n",
            "        (K): Linear(in_features=384, out_features=384, bias=False)\n",
            "        (V): Linear(in_features=384, out_features=384, bias=False)\n",
            "        (ff): Linear(in_features=384, out_features=384, bias=True)\n",
            "        (attn_dropout): Dropout2d(p=0.1, inplace=False)\n",
            "        (ff_dropout): Dropout(p=0.1, inplace=False)\n",
            "      )\n",
            "      (ff): FeedForwardRelativeAttention(\n",
            "        (fc1): Conv2d(384, 1536, kernel_size=(1, 1), stride=(1, 1))\n",
            "        (fc2): Conv2d(1536, 384, kernel_size=(1, 1), stride=(1, 1))\n",
            "      )\n",
            "    )\n",
            "    (3): TransformerWithRelativeAttention(\n",
            "      (norm): LayerNorm((384,), eps=1e-05, elementwise_affine=True)\n",
            "      (attention): RelativeAttention(\n",
            "        (Q): Linear(in_features=384, out_features=384, bias=False)\n",
            "        (K): Linear(in_features=384, out_features=384, bias=False)\n",
            "        (V): Linear(in_features=384, out_features=384, bias=False)\n",
            "        (ff): Linear(in_features=384, out_features=384, bias=True)\n",
            "        (attn_dropout): Dropout2d(p=0.1, inplace=False)\n",
            "        (ff_dropout): Dropout(p=0.1, inplace=False)\n",
            "      )\n",
            "      (ff): FeedForwardRelativeAttention(\n",
            "        (fc1): Conv2d(384, 1536, kernel_size=(1, 1), stride=(1, 1))\n",
            "        (fc2): Conv2d(1536, 384, kernel_size=(1, 1), stride=(1, 1))\n",
            "      )\n",
            "    )\n",
            "    (4): TransformerWithRelativeAttention(\n",
            "      (norm): LayerNorm((384,), eps=1e-05, elementwise_affine=True)\n",
            "      (attention): RelativeAttention(\n",
            "        (Q): Linear(in_features=384, out_features=384, bias=False)\n",
            "        (K): Linear(in_features=384, out_features=384, bias=False)\n",
            "        (V): Linear(in_features=384, out_features=384, bias=False)\n",
            "        (ff): Linear(in_features=384, out_features=384, bias=True)\n",
            "        (attn_dropout): Dropout2d(p=0.1, inplace=False)\n",
            "        (ff_dropout): Dropout(p=0.1, inplace=False)\n",
            "      )\n",
            "      (ff): FeedForwardRelativeAttention(\n",
            "        (fc1): Conv2d(384, 1536, kernel_size=(1, 1), stride=(1, 1))\n",
            "        (fc2): Conv2d(1536, 384, kernel_size=(1, 1), stride=(1, 1))\n",
            "      )\n",
            "    )\n",
            "  )\n",
            "  (s4): Sequential(\n",
            "    (0): TransformerWithRelativeAttention(\n",
            "      (pool): MaxPool2d(kernel_size=(2, 2), stride=(2, 2), padding=0, dilation=1, ceil_mode=False)\n",
            "      (conv): Conv2d(384, 768, kernel_size=(1, 1), stride=(1, 1))\n",
            "      (norm): LayerNorm((384,), eps=1e-05, elementwise_affine=True)\n",
            "      (attention): RelativeAttention(\n",
            "        (Q): Linear(in_features=384, out_features=768, bias=False)\n",
            "        (K): Linear(in_features=384, out_features=768, bias=False)\n",
            "        (V): Linear(in_features=384, out_features=768, bias=False)\n",
            "        (ff): Linear(in_features=768, out_features=768, bias=True)\n",
            "        (attn_dropout): Dropout2d(p=0.1, inplace=False)\n",
            "        (ff_dropout): Dropout(p=0.1, inplace=False)\n",
            "      )\n",
            "      (ff): FeedForwardRelativeAttention(\n",
            "        (fc1): Conv2d(768, 3072, kernel_size=(1, 1), stride=(1, 1))\n",
            "        (fc2): Conv2d(3072, 768, kernel_size=(1, 1), stride=(1, 1))\n",
            "      )\n",
            "    )\n",
            "    (1): TransformerWithRelativeAttention(\n",
            "      (norm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
            "      (attention): RelativeAttention(\n",
            "        (Q): Linear(in_features=768, out_features=768, bias=False)\n",
            "        (K): Linear(in_features=768, out_features=768, bias=False)\n",
            "        (V): Linear(in_features=768, out_features=768, bias=False)\n",
            "        (ff): Linear(in_features=768, out_features=768, bias=True)\n",
            "        (attn_dropout): Dropout2d(p=0.1, inplace=False)\n",
            "        (ff_dropout): Dropout(p=0.1, inplace=False)\n",
            "      )\n",
            "      (ff): FeedForwardRelativeAttention(\n",
            "        (fc1): Conv2d(768, 3072, kernel_size=(1, 1), stride=(1, 1))\n",
            "        (fc2): Conv2d(3072, 768, kernel_size=(1, 1), stride=(1, 1))\n",
            "      )\n",
            "    )\n",
            "  )\n",
            "  (heads): ModuleList(\n",
            "    (0): ProjectionHead(\n",
            "      (fc1): Linear(in_features=768, out_features=768, bias=True)\n",
            "      (act_fn): Mish()\n",
            "      (dropout): Dropout(p=0.1, inplace=False)\n",
            "      (norm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
            "      (fc2): Linear(in_features=768, out_features=100, bias=True)\n",
            "    )\n",
            "  )\n",
            ")\n"
          ]
        }
      ],
      "source": [
        "\n",
        "batch_size = 8\n",
        "num_epochs = 3\n",
        "learning_rate = 0.001\n",
        "num_classes = 100\n",
        "model_config = 'coatnet-0'  \n",
        "\n",
        "\n",
        "input_height = 64\n",
        "input_width = 64\n",
        "\n",
        "\n",
        "model = CoAtNet(input_height, input_width, 3, config=model_config, num_classes=num_classes).to(device)\n",
        "\n",
        "criterion = nn.CrossEntropyLoss()\n",
        "optimizer = optim.Adam(model.parameters(), lr=learning_rate)\n",
        "\n",
        "print(model)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "d1Krl2E_38aB"
      },
      "outputs": [],
      "source": [
        "\n",
        "\n",
        "\n",
        "train_loss_list = []\n",
        "train_acc_list = []\n",
        "\n",
        "\n",
        "for epoch in range(num_epochs):\n",
        "    model.train()\n",
        "    running_loss = 0.0\n",
        "    correct = 0\n",
        "    total = 0\n",
        "\n",
        "    for inputs, labels in trainloader:\n",
        "        inputs, labels = inputs.to(device), labels.to(device)\n",
        "\n",
        "        optimizer.zero_grad()\n",
        "        outputs = model(inputs)\n",
        "        loss = criterion(outputs, labels)\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "\n",
        "        running_loss += loss.item()\n",
        "        _, predicted = outputs.max(1)\n",
        "        total += labels.size(0)\n",
        "        correct += predicted.eq(labels).sum().item()\n",
        "\n",
        "    epoch_loss = running_loss / len(trainloader)\n",
        "    epoch_acc = correct / total\n",
        "    train_loss_list.append(epoch_loss)\n",
        "    train_acc_list.append(epoch_acc)\n",
        "\n",
        "    print(f\"Epoch [{epoch+1}/{num_epochs}]  Loss: {epoch_loss:.4f}  Accuracy: {epoch_acc:.4f}\")\n"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
